import faiss
import numpy as np
import pandas as pd
from exception_hyd.faiss_database_exception import FaissDataBaseException
from numpy.typing import NDArray
from typing import List, Tuple

# 定义一个类型别名
Matrix = List[List[Tuple[str, float]]]


class FaissDataBase:
    """
    自定义一个faiss工具类
    """

    d = 0
    embedding_data_path = ""
    knowledge_data_path = ""
    embedding_data = None
    knowledge_data = None
    filetype_of_knowledge_database = "csv"

    def __init__(self, dimension: int,
                 embedding_data_path: str = "",
                 knowledge_data_path: str = "",
                 filetype_of_kb: str = "csv"):

        # 判断dimension, embedding_data_path, knowledge_data_path值是否为空
        if not (dimension and embedding_data_path and knowledge_data_path):
            raise FaissDataBaseException("The existence parameter is empty")

        # 判断dimension, embedding_data_path, knowledge_data_path的类型是否正确
        if not (isinstance(dimension, int) and
                isinstance(embedding_data_path, str) and
                isinstance(knowledge_data_path, str)):
            raise FaissDataBaseException("There is an incorrect parameter type")

        # 训练矩阵的维度
        self.d = dimension

        # 向量数据的文件路径
        self.embedding_data_path = embedding_data_path

        # 知识库文件的路径
        self.knowledge_data_path = knowledge_data_path

        # 加载知识库
        self.embedding_data = np.load(self.embedding_data_path)

        # 获取知识库的文件类型
        self.filetype_of_knowledge_database = filetype_of_kb

        # 加载知识库
        if self.filetype_of_knowledge_database == "csv":
            try:
                self.knowledge_data = pd.read_csv(self.knowledge_data_path, encoding="utf-8")
            except UnicodeDecodeError:
                self.knowledge_data = pd.read_csv(self.knowledge_data_path, encoding="gbk")

        elif self.filetype_of_knowledge_database == "xlsx":
            self.knowledge_data = pd.read_excel(self.knowledge_data_path)

        else:
            raise FaissDataBaseException("File type not recognized, only the file type of csv and xlsx are supported")

    def create_index_l2(self) -> faiss.IndexFlatL2:
        """
        创建一个L2索引

        :return: faiss.IndexFlatL2对象
        """

        l2_index = faiss.IndexFlatL2(self.d)
        l2_index.add(self.embedding_data)
        return l2_index

    def search_vector_text(self, faiss_obj: faiss.IndexFlatL2,
                           top_k: int,
                           embedding_data: NDArray[np.float32],
                           threshold_value: float = -1) -> Matrix:
        """
        查询向量知识库并返回相关的知识条例

        :param faiss_obj: 获取一个索引
        :param top_k: 获取相关的top_k个知识条例
        :param embedding_data: 需要查询的向量数据
        :param threshold_value: 阀值
        :return:
        """

        distance_point_matrix, index_matrix = faiss_obj.search(embedding_data, top_k)
        data_matrix = []

        # 当没有设置阀值时
        if threshold_value == -1:

            # 获取知识条例
            for i, k in enumerate(index_matrix):
                text_data = self.knowledge_data.iloc[k]["input"]
                confidence_coefficient = distance_point_matrix[i]
                tuple_data = list(zip(text_data, confidence_coefficient))
                data_matrix.append(tuple_data)

        # 当设置了阀值时
        else:

            # 储存满足阀值条件的索引值
            satisfy_index = []
            for i in range(len(distance_point_matrix)):
                row_i = []
                for j in range(len(distance_point_matrix[i])):
                    if distance_point_matrix[i][j] < threshold_value:
                        index_tuple = (i, j)
                        row_i.append(index_tuple)
                satisfy_index.append(row_i)

            not_empty = [i for i in satisfy_index if len(i) != 0]
            if not_empty:

                # 储存满足阀值条件的向量库的索引值
                index_matrix_ = []

                # 储存满足条件的向量库数据对应的distance, 即置信度
                distance_point_matrix_ = []
                for i in satisfy_index:
                    row_i_index = []
                    row_i_distance = []
                    for j in i:
                        row_i_index.append(index_matrix[j[0]][j[1]])
                        row_i_distance.append(distance_point_matrix[j[0]][j[1]])
                    index_matrix_.append(row_i_index)
                    distance_point_matrix_.append(row_i_distance)

                # 获取知识条例
                for i, k in enumerate(index_matrix_):
                    text_data = self.knowledge_data.iloc[k]["input"]
                    confidence_coefficient = distance_point_matrix_[i]
                    tuple_data = list(zip(text_data, confidence_coefficient))
                    data_matrix.append(tuple_data)

        return data_matrix
